import os
from TrainData import TrainData

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

from keras import Sequential
from keras import callbacks
from keras.layers import Dense

from keras2cpp import export_model

class DNNTrainer:
    
    ## 数据对象
    __trainObj = TrainData("")
    
    ## 神经元数量
    __neuronsCount = 24
    ## 隐藏层数量
    __hiddenLayersCount = 2
    ## 优化器
    __optimizer = 'Adam'
    ## 激活函数
    __activation = 'relu'
    ## 验证集比例
    __validationRatio = 0
    
    ## 训练历史
    __trainHistory = callbacks.History()
    ## 早停机制
    __earlyStop = callbacks.EarlyStopping(patience = 10000)
    ## 预测模型
    __model = Sequential()

    
    ## Constructor to initial dnnTrainer
    ## trainObj:    TrainData Object
    ## neCount:     神经元数量
    ## hlCount:     隐藏层数量
    ## optimizer:   优化器
    ## activation:  激活函数
    def __init__(self, trainObj, neCount, hlCount, activation, optimizer):
        self.__trainObj = trainObj
        self.__neuronsCount = neCount
        self.__hiddenLayersCount = hlCount
        self.__optimizer = optimizer
        self.__activation = activation
        
    def SetEarlyStop(self, earlyPatience, monitor ='val_mse', min_delta= 0.0002):
        ##早停机制 验证集MSE:val_mse  
        self.__earlyStop = callbacks.EarlyStopping(monitor = monitor, 
                                          min_delta = min_delta, patience = earlyPatience)
    
    ## 数据训练
    def Train(self, epochs, validationRatio = 0.2, loss = 'mse', verbose = 1):
        ## 验证集数量
        self.__validationRatio = validationRatio
         ## 输入、输出参数
        inValues = self.__trainObj.GetTrainValues(True)
        outValues = self.__trainObj.GetTrainValues(False)
        ## DNN model
        self.__model = self.__getDNNModel(inValues.shape[1], outValues.shape[1])
        ##损失函数 优化方法 
        self.__model.compile(loss = loss, optimizer= self.__optimizer,
                      metrics=['mae','mse'])
        ##训练 verbose:输出方式枚举（0，1，2）
        self.__trainHistory = self.__model.fit(inValues, outValues, epochs = epochs, 
                    verbose = verbose, validation_split = self.__validationRatio,
                    callbacks=[self.__earlyStop])
        
    ## 绘制损失函数   
    def PlotLoss(self):
        ## 获得训练历史
        hist = pd.DataFrame(self.__trainHistory.history)
        hist['epoch'] = self.__trainHistory.epoch
        ## 创建绘图器
        plt.figure()
        ## 横轴纵轴名称
        plt.xlabel('Epoch')
        plt.ylabel('MSE')
        ## 训练过程损失函数
        plt.plot(hist['epoch'],hist['loss'], label='Train Error')
        ## 判断是否有验证集
        if self.__validationRatio > 0.01 :
            plt.plot(hist['epoch'],hist['val_loss'], label='val error')
        plt.legend()
        plt.show()
    
    ## 绘制预测结果散点图
    def PlotPrediction(self, isTrain, plotIndex): 
        ## 获得输入输出参数
        inValues = self.__trainObj.GetTrainValues(True) if isTrain else self.__trainObj.GetTestValues(True)
        outValues = self.__trainObj.GetTrainValues(False) if isTrain else self.__trainObj.GetTestValues(False)         
        outFeatures = self.__trainObj.GetTrainFeatures(False)
        ## 绘图
        self. __plotPrediction(inValues, outValues, outFeatures, plotIndex)
        
    ## 绘制预测结果散点图
    def PlotPredictions(self, isTrain): 
        ## 获得输入输出参数
        inValues = self.__trainObj.GetTrainValues(True) if isTrain else self.__trainObj.GetTestValues(True)
        outValues = self.__trainObj.GetTrainValues(False) if isTrain else self.__trainObj.GetTestValues(False)         
        outFeatures = self.__trainObj.GetFeatures(False)
        ## 遍历特性 绘图
        for plotIndex in range(0, outFeatures.size):
            self. __plotPrediction(inValues, outValues, outFeatures, plotIndex)
        
     ## 绘制预测结果散点图
    def __plotPrediction(self, inValues, outValues, outFeatures, plotIndex):
        ## 预测结果
        predictions = self.__model.predict(inValues)
        ## 绘图
        plt.scatter(outValues[:,plotIndex], predictions[:,plotIndex])
        plt.xlabel('True Values')
        plt.ylabel('Predictions')
        plt.axis('equal')
        plt.axis('square')
        ## 最大值
        outValueMax = outValues[:,plotIndex].max()
        ## 坐标轴限定
        plt.xlim([0,outValueMax])
        plt.ylim([0,outValueMax])
        plt.plot([0, outValueMax], [0, outValueMax], label = outFeatures[plotIndex])
        plt.legend()
        plt.show()
            
    def SaveModel(self, directPath, exportModelName ='M'):
        ## 创建文件夹
        self.__trainObj.CreateDirectory(directPath)
        ## 输入、输出参数
        intrainValues = self.__trainObj.GetTrainValues(True)
        outrainValues = self.__trainObj.GetTrainValues(False)
        intestValues = self.__trainObj.GetTestValues(True)
        outestValues = self.__trainObj.GetTestValues(False)
        ## 预测
        trainPredicts = self.__model.predict(intrainValues)
        testPredicts = self.__model.predict(intestValues)
        ## 输出
        self.__mergeSave(intrainValues, trainPredicts, directPath + '\\trainPredict.txt')
        self.__mergeSave(intestValues, testPredicts, directPath + '\\testPredict.txt')
        self.__mergeSave(intrainValues, outrainValues, directPath + '\\trainData.txt')
        self.__mergeSave(intestValues, outestValues, directPath + '\\testData.txt')
        ## 加载历史
        hist = pd.DataFrame(self.__trainHistory.history)
        hist['epoch'] = self.__trainHistory.epoch
        hist.to_csv(directPath + '\\history.txt', index = False)
        ## 模型固化
        export_model(self.__model, directPath + '\\' + exportModelName + '.model')
        self.__model.save(directPath + '\\' + exportModelName + ".h5")
        
    
    ## 创建深度学习模型
    def __getDNNModel(self, inValueSize, outValueSize):
        ## 模型初始化
        model = Sequential()
        ## 第一隐藏层
        model.add(Dense(self.__neuronsCount, activation = self.__activation, 
                        input_dim = inValueSize))
        ## 其他隐藏层
        for i in range(1, self.__hiddenLayersCount):
            model.add(Dense(self.__neuronsCount, activation = self.__activation))
        ## 输出层
        model.add(Dense(outValueSize))
        ## 模型结构查看
        model.summary()
        return model
    
    ## 预测并返回输入参数
    def __mergeSave(self, inValues, outValues, filePath):
        ## 列表合并 
        mergeList = np.hstack((inValues, outValues))
        ## 属性名称
        features = self.__trainObj.GetAllFeatures()
        ## 合并
        mergeDF = pd.DataFrame(mergeList,columns = features)
        ## 输出
        mergeDF.to_csv(filePath, index = False)